from langchain import hub
from customize.get_ollama import GetOllama
from langgraph.prebuilt import create_react_agent
from customize.bocha_web_search import get_bocha_tool
from typing import Literal
from langgraph.graph import END
from typing import Union
import operator
from typing import Annotated, List, Tuple
from typing_extensions import TypedDict
from pydantic import BaseModel, Field
from langgraph.graph import StateGraph, START
from langchain_core.prompts import ChatPromptTemplate
from customize.save_image import save_graph
import asyncio


prompt = hub.pull("ih/ih-react-agent-executor")
prompt.pretty_print()
tools = [get_bocha_tool()]
llm = GetOllama(ip=GetOllama.ailab_linux_ip, model_name="qwen2.5:14b", model_type=1)()
agent_executor = create_react_agent(llm, tools, state_modifier=prompt)


class PlanExecute(TypedDict):
    input: str
    plan: List[str]
    past_steps: Annotated[List[Tuple], operator.add]
    response: str


class Plan(BaseModel):
    """Plan to follow in future"""

    steps: List[str] = Field(
        description="different steps to follow, should be in sorted order"
    )


planner_prompt = ChatPromptTemplate.from_messages(
    [
        (
            "system",
#             """For the given objective, come up with a simple step by step plan. \
# This plan should involve individual tasks, that if executed correctly will yield the correct answer. Do not add any superfluous steps. \
# The result of the final step should be the final answer. Make sure that each step has all the information needed - do not skip steps.""",
            """
            对于给定的目标，提出一个简单的分步计划。
            此计划应涉及单个任务，如果正确执行将产生正确的答案。不要添加任何多余的步骤。
            最后一步的结果应该是最终答案。确保每个步骤都包含所需的所有信息 —— 不要跳过步骤。
            请按下面的格式输出计划：
                "steps":["计划1","计划2","计划3"]
"""
        ),
        ("placeholder", "{messages}"),
    ]
)
planner_llm = GetOllama(ip=GetOllama.ailab_linux_ip, model_name="qwen2.5:14b", model_type=1, temperature=0)()
planner = planner_prompt | planner_llm.with_structured_output(Plan)

result = planner.invoke(
    {
        "messages": [
            ("user", "今年巴黎奥运会乒乓球男单冠军，他有没有结婚?")
        ]
    }
)


print(result)